package org.daochong.web;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.io.UnsupportedEncodingException;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;

public class HttpResponseWrapper extends HttpServletResponseWrapper {

	public static String DEFAULT_ENCODING = "UTF-8";

	private OutputStream out;
	private PrintWriter writer;
	private ServletOutputStream servletOut;
	private byte[] data;

	public HttpResponseWrapper(HttpServletResponse response) {
		super(response);
		out = new ByteArrayOutputStream();
	}

	@Override
	public ServletOutputStream getOutputStream() throws IOException {
		if (this.servletOut == null) {
			this.servletOut = new ResponseServletOutputStream(getResponse().getOutputStream());
		}
		return servletOut;
	}

	@Override
	public PrintWriter getWriter() throws IOException {
		if (this.writer == null) {
			this.writer = new ResponsePrintWriter(this.getOutputStream(),
					(this.getCharacterEncoding() == null ? DEFAULT_ENCODING : this.getCharacterEncoding()));
		}
		return this.writer;
	}

	public String getContent() {
		String c = (this.getCharacterEncoding() != null ? this.getCharacterEncoding() : DEFAULT_ENCODING);
		try {
			if (this.data != null) {
				return new String(this.data, c);
			}
			return ((ByteArrayOutputStream) this.out).toString(c);
		} catch (UnsupportedEncodingException e) {
			throw new RuntimeException(e);
		}
	}

	public byte[] getContenData() {
		if (this.data != null) {
			return this.data;
		}
		return ((ByteArrayOutputStream) this.out).toByteArray();
	}

	private class ResponseServletOutputStream extends ServletOutputStream {

		private ServletOutputStream so;

		public ResponseServletOutputStream(ServletOutputStream out) {
			this.so = out;
		}

		@Override
		public void flush() throws IOException {
			out.flush();
			this.so.flush();
		}

		@Override
		public void close() throws IOException {
			data = ((ByteArrayOutputStream) out).toByteArray();
			out.close();
			this.so.close();
		}

		@Override
		public void write(int b) throws IOException {
			out.write(b);
			this.so.write(b);
		}

		@Override
		public void write(byte[] b, int off, int len) throws IOException {
			out.write(b, off, len);
			this.so.write(b,off,len);
		}
	}

	private class ResponsePrintWriter extends PrintWriter {

		public ResponsePrintWriter(OutputStream out, String characterEncoding) throws UnsupportedEncodingException {
			super(new OutputStreamWriter(out, characterEncoding));
		}

		@Override
		public void write(char buf[], int off, int len) {
			super.write(buf, off, len);
			super.flush();
		}

		@Override
		public void write(String s, int off, int len) {
			super.write(s, off, len);
			super.flush();
		}

		@Override
		public void write(int c) {
			super.write(c);
			super.flush();
		}

	}
}
